/*
 *  emu3d
 *
 *	Code for merging 2D diffraction images into a single 3D volume. 
 *  Created by Anton Barty on 4/9/10.
 *  Copyright 2010 all rights reserved.
 *  
 *
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>
#include <hdf5.h>
#include <math.h>
#include <pthread.h>
#include "phaseit3d.h"


void testcode(void) {
	
	
	
	/*
	 *	Test reading of 2D HDF5 files
	 *
		cData2d		test;
		test.readHDF5((char*) "test.h5");
		printf("nn=%li\n",test.nn);
		printf("nx,ny=%li,%li\n",test.nx, test.ny);
		for (long i=0; i<test.nn; i++) {
			printf("%f\t",test.data[i]);
		}
		printf("\n");
	 */

	
	/*
	 *	Test reading and writing 3D HDF5 files
	 *
		cData3d		test3;
		test3.readHDF5((char*)"test3.h5");
		test3.writeHDF5((char*)"test3a.h5");
		exit(1);
		*/
	
	H5close();
}


int main(int argc, char **argv)
{
	
	char filename[1024];

	
	/*
	 *	System checks
	 */
	printf("\nChecking system\n");
	printf("\tsizeof(char) = %li\n",sizeof(char));
	printf("\tsizeof(short int) = %li\n",sizeof(short int));
	printf("\tsizeof(int) = %li\n",sizeof(int));
	printf("\tsizeof(long) = %li\n",sizeof(long));
	printf("\tsizeof(float) = %li\n",sizeof(float));
	printf("\tsizeof(double) = %li\n",sizeof(double));
	printf("\tsizeof(fftwf_complex) = %li\n",sizeof(fftwf_complex));
	printf("\tsizeof(fftw_complex) = %li\n",sizeof(fftw_complex));
	printf("\n");




	/*
	 *	Process command line arguments (if any)
	 */
	cSettings	settings;
	settings.setDefaults();
	settings.parseCommandLineArguments(argc, argv);

	
	/*
	 *	Global variables
	 */
	cGlobal		global;
	global.init(&settings);
	
	
	/*
	 *	FFTW testing
	 */
	if(settings.fftwSpeedTest) {
		printf("\n");
		printf("Performing FFTW speed test\n");
		printf("\tProblem size is: %li^3\n",settings.fftwSpeedTestN); 

		cFFTWdata	test;
		test.create(settings.fftwSpeedTestN);
		test.createFFTWplan(settings.nthreads);
		test.speedTest(10);
		printf("\tSpeed test has finished...\n");
		
		exit(1);
	}
	

	/*
	 *	Read in diffraction volume
	 */
	cData3d		diffractionVolume;
	printf("\n");
	printf("Reading data volume: %s\n",settings.diffractionDataFile);
	diffractionVolume.readHDF5(settings.diffractionDataFile);
	printf("\tProblem size is: %li^3\n",diffractionVolume.nx); 
	
	if(!settings.dvAmplitude) {
		printf("\tConverting from intensity to amplitude\n"); 
		for(long p=0; p<diffractionVolume.nn; p++) {
			if(diffractionVolume.data[p] >= 0 && diffractionVolume.data[p] != settings.dvMissingData )
				diffractionVolume.data[p] = sqrt(diffractionVolume.data[p]);
			else
				diffractionVolume.data[p] = settings.dvMissingData;
		}
	}
	if(!settings.dvZeroInCorner) {
		printf("\tShifting diffraction volume center into corner\n"); 
		diffractionVolume.shift(diffractionVolume.nx/2, diffractionVolume.ny/2, diffractionVolume.nz/2);
	}

	
	/*
	 *	Copy data into fft array
	 */
	cFFTWdata	model;
	model.create(diffractionVolume.nx);
	model.createFFTWplan(settings.nthreads);
	
	for(long p=0; p<model.nn; p++) {
		if(diffractionVolume.data[p] != settings.dvMissingData) {
			model.data[p][0] = diffractionVolume.data[p];
			model.data[p][1] = 0.0;		
		}
		else {
			model.data[p][0] = 0;
			model.data[p][1] = 0.0;		
		}
	}
	
	
	/*
	 *	FFTW sanity check - to make sure FFTs do not screw up the arrays
	 */
	if(settings.fftwSanityCheck) {
		printf("\n");
		printf("Performing FFTW sanity checks\n");
		printf("\tProblem size is: %li^2\n",settings.fftwSpeedTestN); 
		model.fft(FFTW_BACKWARD);
		printf("\tFFT sanity check finished\n");
		
		exit(1);
	}
	
	
	/*
	 *	Initial support
	 */
	cData3d		supportVolume;
	
	if(strcmp(settings.supportDataFile, "none")) {	
		// Read initial support from file
		printf("\n");
		printf("Reading initial support volume: %s\n",settings.supportDataFile);
		supportVolume.readHDF5(settings.supportDataFile);
		
		if(supportVolume.nx != diffractionVolume.nx || supportVolume.ny != diffractionVolume.ny || supportVolume.nz != diffractionVolume.nz) {
			printf("\tDimensions of support and diffraction volume do not agree\n"); 
			printf("\tAborting...\n");
			exit(1);
		}
		
		if(!settings.supportInCorner) {
			printf("\tShifting support volume center into corner\n"); 
			supportVolume.shift(supportVolume.nx/2, supportVolume.ny/2, supportVolume.nz/2);
		}
	}
	else {
		// Estimate support from autocorrelation
		printf("\n");
		printf("Creating initial support volume from autocorrelation\n");
		model.fft(FFTW_BACKWARD);
		updateSupport(&model, &settings, 10, settings.swThresh0);
		supportVolume.create(model.nx);
		for(long p=0; p<model.nn; p++) {
			supportVolume.data[p] = model.data[p][0];
		}
	}
	
	
	/*
	 *	Initial object guess (for now, only real-space)
	 */
	cData3d		guessVolume;
	
	if(strcmp(settings.guessDataFile, "none")) {	
		// Read initial support from file
		printf("\n");
		printf("Reading initial guess: %s\n",settings.guessDataFile);
		guessVolume.readHDF5(settings.guessDataFile);
		
		if(guessVolume.nx != diffractionVolume.nx || guessVolume.ny != diffractionVolume.ny || guessVolume.nz != diffractionVolume.nz) {
			printf("\tDimensions of guess and diffraction volume do not agree\n"); 
			printf("\tAborting...\n");
			exit(1);
		}
		
		if(!settings.guessInCorner) {
			printf("\tShifting support volume center into corner\n"); 
			guessVolume.shift(diffractionVolume.nx/2, diffractionVolume.ny/2, diffractionVolume.nz/2);
		}
	}
	else {
		// Default is 1 where support is defined
		guessVolume.create(model.nx);
		for(long p=0; p<model.nn; p++) {
			if (supportVolume.data[p] == 0)
				guessVolume.data[p] = 0;				
			else
				guessVolume.data[p] = 1;	
		}
	}

	
	// Do some sanity checks on the support and guess
	sprintf(filename,"%s/support00000.h5",settings.outdir);
	printf("Writing initial support: %s\n", filename);
	supportVolume.shift(supportVolume.nx/2, supportVolume.ny/2, supportVolume.nz/2);
	supportVolume.writeHDF5((char*)filename);
	supportVolume.shift(supportVolume.nx/2, supportVolume.ny/2, supportVolume.nz/2);

	
	// Stop here for now
	//printf("clean exit\n");
	//exit(1);

	
	/*
	 *	Call the phase retrieval function that does all the work
	 */
	printf("Calling phaser\n");
	phaseit(&diffractionVolume, &supportVolume, &guessVolume, &settings);
	
	
	
	
	
	/*
	 *	Exit cleanly
	 */
	printf("Clean exit\n");

     
}
